Proposed changes

Fixes a CPU copy bug where copy_general_general iterated over src.size() even when callers provided a smaller data_shape.

DynamicSlice::eval_cpu passes the output shape into copy_cpu_inplace, so CPU dynamic slices with higher-rank collapsed shapes could copy past the destination bounds. In practice this caused crashes and unnecessary extra work. The fix computes the copy size from data_shape, matching the requested copy region.

Reproduce

Before this patch, the CPU path crashes while the GPU path returns the expected slice.

import mlx.core as mx

x = mx.arange(5 * 6 * 7 * 8).reshape(5, 6, 7, 8)
starts = mx.array([1, 2, 3, 4])
axes = (0, 1, 2, 3)
sizes = (2, 2, 2, 2)

with mx.stream(mx.gpu):
    expected = x[1:3, 2:4, 3:5, 4:6]
    gpu_out = mx.slice(x, starts, axes, sizes)
    print("gpu ok:", mx.array_equal(expected, gpu_out).item())

with mx.stream(mx.cpu):
    expected = x[1:3, 2:4, 3:5, 4:6]
    cpu_out = mx.slice(x, starts, axes, sizes)
    print("cpu ok:", mx.array_equal(expected, cpu_out).item())

Behavior:

• GPU: prints gpu ok: True
• CPU: segfaults on the dynamic slice

Tests

python -m pytest python/tests/test_ops.py -k dynamic_slicing -q
DEVICE=cpu ./build/tests/tests --test-case='test dynamic slice'

Checklist

Put an x in the boxes that apply.

• I have read the CONTRIBUTING document
• I have run pre-commit run --all-files to format my code / installed pre-commit prior to committing changes
• I have added tests that prove my fix is effective or that my feature works